###############################################################################
# 功能：
#       yolo推理class封装
# 说明：
#
# 作者：
#       罗培元
# 日期：
#       2018 - 12 - 8
# 修改人：
#       xxx
#       xxx
# 重要修改说明：
#
###############################################################################
import os, sys
import argparse
import json
import cv2
from PIL import Image

sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "./../"))

from pylab import *

sys.path.append(os.path.dirname(__file__))

from utils.utils import get_yolo_boxes, makedirs
from utils.bbox import draw_boxes, get_boxes
from keras.models import load_model
from keras import backend as K
# from tqdm import tqdm
import numpy as np
import time

# add by suyongsheng
import tensorflow as tf
from multiprocessing import cpu_count


# anchor for [4,6, 8,12, 15,20, 24,33, 35,51, 54,65, 55,115, 91,111, 115,215]
# old anchor [55,69, 75,234, 133,240, 136,129, 142,363, 203,290, 228,184, 285,359, 341,260]
class Yolo:
    def __init__(self, config_path):
        self.set_context()
        with open(config_path) as config_buffer:
            self.config = json.load(config_buffer)

        self.net_h, self.net_w = 416, 416
        # self.net_h, self.net_w = 288, 288
        self.obj_thresh, self.nms_thresh = 0.5, 0.45
        os.environ['CUDA_VISIBLE_DEVICES'] = self.config['train']['gpus']
        self.infer_model = load_model(self.config['train']['saved_weights_name'])

    def set_context(self, num_cores=3, GPU=True,
                    CPU=False):
        if GPU:
            num_GPU = 1
            num_CPU = 1
        if CPU:
            num_CPU = 1
            num_GPU = 0
        config = tf.ConfigProto(intra_op_parallelism_threads=num_cores,
                                inter_op_parallelism_threads=num_cores, allow_soft_placement=True,
                                device_count={'CPU': num_CPU, 'GPU': num_GPU})
        session = tf.Session(config=config)
        K.set_session(session)

    def predict(self, image, from_file=False, show_im=False, verbose=True):
        if from_file:
            image = cv2.imread(image)
        time_start = time.time()
        # predict the bounding boxes
        boxes = \
            get_yolo_boxes(self.infer_model, [image], self.net_h, self.net_w, self.config['model']['anchors'],
                           self.obj_thresh, self.nms_thresh, verbose=verbose)[0]
        time_end = time.time()
        print('yolo predict cost time=', time_end - time_start)
        xyc_pack = get_boxes(image, boxes, self.config['model']['labels'], self.obj_thresh)
        if show_im:
            image = draw_boxes(image, boxes, self.config['model']['labels'], self.obj_thresh)
            imshow(image)
            show()
        return xyc_pack


def case1():
    model_yolo = Yolo("./config_ut.json")
    image = cv2.imread("../test_im1080.jpg")
    model_yolo.predict(image, show_im=False)
    model_yolo.predict(image, show_im=False)
    model_yolo.predict(image, show_im=False)
    model_yolo.predict(image, show_im=False)
    model_yolo.predict(image, show_im=True)


if __name__ == '__main__':
    case1()
